import torch
from torch import nn
from __00__config import Config

config = Config()
class BertClassifierModel(nn.Module):
	def __init__(self):
		super().__init__()
		# 获取预训练模型
		self.bert_model = config.bert_model
		# 添加一层线性层
		self.linear = nn.Linear(config.bert_config.hidden_size, config.num_classes)

	def forward(self, input_ids, attention_mask):
		outputs = self.bert_model(
			input_ids=input_ids,  # 输入的token ID
			attention_mask=attention_mask  # 注意力掩码用于区分有效token和填充token
		)
		logits = self.linear(outputs['pooler_output'])
		return logits


if __name__ == '__main__':
	tokenizer = config.bert_tokenizer

	# 示例文本
	texts = ["王者荣耀", "今天天气真好"]

	# 编码文本
	encoded_inputs = tokenizer(
		texts,
		# padding=True,  #  所有的填充到文本最大长度
		padding="max_length",  # 所有的填充到指定的max_length长度
		max_length=6,
		truncation=True,  # 如果超出指定的max_length长度，则截断
		return_tensors="pt"  # 返回 pytorch 张量,"pt" 时，分词器会将输入文本转换为模型可接受的格式
	)

	# 获取 input_ids 和 attention_mask
	input_ids = encoded_inputs["input_ids"]
	attention_mask = encoded_inputs["attention_mask"]
	print('input_ids:', input_ids)
	print('attention_mask:', attention_mask)
	print('==================================================================')
	# 创建自定义的bert模型
	model = BertClassifierModel()  # __init__()执行了
	# 预测
	logits = model(input_ids, attention_mask)  # forward()执行了
	print(logits)  # 每一行对应一个样本，每个数字表示该样本属于某一类别的“得分”（logit），没有经过 softmax 归一化。
	print('-------------------------------')
	# 获取预测概率
	probs = torch.softmax(logits, dim=-1)
	print(probs)  # 归一化后该样本属于某类的概率（范围在 0~1 之间）,概率最高的就是预测结果
	print('-------------------------------')
	# 取预测结果
	preds = torch.argmax(logits, dim=-1)
	print(preds)  # 得到每个样本的预测类别。表示两个输入文本被模型预测为类别 6（从 0 开始计数）。
